"""
增强版协同编辑服务
实现更完善的Operational Transformation(OT)算法和CRDT支持
"""

import json
import asyncio
from typing import Dict, List, Any, Optional, Set
from datetime import datetime
from fastapi import WebSocket, WebSocketDisconnect
from src.utils.logging import get_logger

logger = get_logger(__name__)


class Operation:
    """操作对象"""
    def __init__(self, op_type: str, position: int = 0, text: str = "", length: int = 0, 
                 user_id: str = "", timestamp: str = "", revision: int = 0):
        self.type = op_type
        self.position = position
        self.text = text
        self.length = length
        self.user_id = user_id
        self.timestamp = timestamp or datetime.now().isoformat()
        self.revision = revision

    def to_dict(self):
        return {
            "type": self.type,
            "position": self.position,
            "text": self.text,
            "length": self.length,
            "user_id": self.user_id,
            "timestamp": self.timestamp,
            "revision": self.revision
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]):
        return cls(
            op_type=data.get("type", ""),
            position=data.get("position", 0),
            text=data.get("text", ""),
            length=data.get("length", 0),
            user_id=data.get("user_id", ""),
            timestamp=data.get("timestamp", ""),
            revision=data.get("revision", 0)
        )


class EnhancedCollaborationManager:
    """增强版协同编辑管理器"""
    
    def __init__(self):
        # 存储文档内容: {document_id: content}
        self.documents: Dict[str, str] = {}
        # 存储文档版本历史: {document_id: {revision: content}}
        self.document_versions: Dict[str, Dict[int, str]] = {}
        # 存储操作历史: {document_id: [operations]}
        self.operation_history: Dict[str, List[Operation]] = {}
        # 存储文档当前版本号: {document_id: revision}
        self.document_revisions: Dict[str, int] = {}
        # 存储用户光标位置: {document_id: {user_id: cursor_position}}
        self.user_cursors: Dict[str, Dict[str, Dict[str, Any]]] = {}
        # 存储文档锁: {document_id: {user_id: lock_info}}
        self.document_locks: Dict[str, Dict[str, Dict[str, Any]]] = {}
        # 存储在线用户: {document_id: {user_id: user_info}}
        self.online_users: Dict[str, Dict[str, Dict[str, Any]]] = {}
    
    def initialize_document(self, document_id: str, initial_content: str = ""):
        """初始化文档"""
        if document_id not in self.documents:
            self.documents[document_id] = initial_content
            self.document_versions[document_id] = {0: initial_content}
            self.operation_history[document_id] = []
            self.document_revisions[document_id] = 0
            self.user_cursors[document_id] = {}
            self.document_locks[document_id] = {}
            self.online_users[document_id] = {}
            logger.info(f"文档 {document_id} 初始化完成")
    
    def apply_operation(self, document_id: str, operation: Operation) -> bool:
        """应用操作到文档"""
        try:
            if document_id not in self.documents:
                self.initialize_document(document_id)
            
            current_content = self.documents[document_id]
            current_revision = self.document_revisions[document_id]
            
            # 检查操作版本是否匹配
            if operation.revision != current_revision:
                logger.warning(f"操作版本不匹配: 期望 {current_revision}, 实际 {operation.revision}")
                return False
            
            # 应用操作
            if operation.type == "insert":
                new_content = current_content[:operation.position] + operation.text + current_content[operation.position:]
            elif operation.type == "delete":
                new_content = current_content[:operation.position] + current_content[operation.position + operation.length:]
            else:
                logger.warning(f"未知操作类型: {operation.type}")
                return False
            
            # 更新文档内容和版本
            new_revision = current_revision + 1
            self.documents[document_id] = new_content
            self.document_versions[document_id][new_revision] = new_content
            self.document_revisions[document_id] = new_revision
            
            # 记录操作历史
            operation.revision = new_revision
            self.operation_history[document_id].append(operation)
            
            logger.info(f"文档 {document_id} 应用操作成功，新版本: {new_revision}")
            return True
        except Exception as e:
            logger.error(f"应用操作失败: {str(e)}")
            return False
    
    def transform_operation(self, document_id: str, op1: Operation, op2: Operation) -> Operation:
        """操作转换算法 (OT)"""
        # 简化的OT实现 - 实际应用中需要更复杂的算法
        if op1.position <= op2.position:
            if op2.type == "insert":
                op1.position += len(op2.text)
            elif op2.type == "delete":
                op1.position -= min(op2.length, op1.position)
        else:
            if op1.type == "insert":
                op2.position += len(op1.text)
            elif op1.type == "delete":
                op2.position -= min(op1.length, op2.position)
        
        return op1
    
    def get_document_content(self, document_id: str) -> Optional[str]:
        """获取文档内容"""
        return self.documents.get(document_id)
    
    def get_document_revision(self, document_id: str) -> int:
        """获取文档版本号"""
        return self.document_revisions.get(document_id, 0)
    
    def get_operation_history(self, document_id: str, limit: int = 100) -> List[Dict[str, Any]]:
        """获取操作历史"""
        if document_id in self.operation_history:
            ops = self.operation_history[document_id][-limit:]
            return [op.to_dict() for op in ops]
        return []
    
    def update_user_cursor(self, document_id: str, user_id: str, cursor_data: Dict[str, Any]):
        """更新用户光标位置"""
        if document_id not in self.user_cursors:
            self.user_cursors[document_id] = {}
        self.user_cursors[document_id][user_id] = cursor_data
    
    def get_user_cursors(self, document_id: str) -> Dict[str, Dict[str, Any]]:
        """获取文档中所有用户的光标位置"""
        return self.user_cursors.get(document_id, {})
    
    def acquire_lock(self, document_id: str, user_id: str, lock_info: Dict[str, Any]) -> bool:
        """获取文档锁"""
        if document_id not in self.document_locks:
            self.document_locks[document_id] = {}
        
        # 检查是否已被锁定
        if document_id in self.document_locks and self.document_locks[document_id]:
            existing_lock = next(iter(self.document_locks[document_id].values()))
            # 如果是同一个用户，允许更新锁
            if existing_lock.get("user_id") != user_id:
                return False
        
        lock_info["user_id"] = user_id
        lock_info["timestamp"] = datetime.now().isoformat()
        self.document_locks[document_id][user_id] = lock_info
        return True
    
    def release_lock(self, document_id: str, user_id: str):
        """释放文档锁"""
        if document_id in self.document_locks and user_id in self.document_locks[document_id]:
            del self.document_locks[document_id][user_id]
            # 如果没有其他锁，清空该文档的锁
            if not self.document_locks[document_id]:
                del self.document_locks[document_id]
    
    def get_lock_info(self, document_id: str) -> Optional[Dict[str, Any]]:
        """获取文档锁信息"""
        if document_id in self.document_locks and self.document_locks[document_id]:
            return next(iter(self.document_locks[document_id].values()))
        return None
    
    def add_user(self, document_id: str, user_id: str, user_info: Dict[str, Any]):
        """添加在线用户"""
        if document_id not in self.online_users:
            self.online_users[document_id] = {}
        self.online_users[document_id][user_id] = user_info
    
    def remove_user(self, document_id: str, user_id: str):
        """移除在线用户"""
        if document_id in self.online_users and user_id in self.online_users[document_id]:
            del self.online_users[document_id][user_id]
            # 同时移除用户的光标和锁
            if document_id in self.user_cursors and user_id in self.user_cursors[document_id]:
                del self.user_cursors[document_id][user_id]
            self.release_lock(document_id, user_id)
    
    def get_online_users(self, document_id: str) -> Dict[str, Dict[str, Any]]:
        """获取在线用户列表"""
        return self.online_users.get(document_id, {})


class EnhancedCollaborationService:
    """增强版协同编辑服务"""
    
    def __init__(self):
        self.manager = EnhancedCollaborationManager()
        self.active_connections: Dict[str, Dict[str, WebSocket]] = {}
    
    async def connect(self, websocket: WebSocket, document_id: str, user_id: str, user_info: Dict[str, Any]):
        """建立WebSocket连接"""
        await websocket.accept()
        
        # 初始化文档
        self.manager.initialize_document(document_id)
        
        # 存储连接
        if document_id not in self.active_connections:
            self.active_connections[document_id] = {}
        self.active_connections[document_id][user_id] = websocket
        
        # 添加在线用户
        self.manager.add_user(document_id, user_id, user_info)
        
        # 发送初始文档状态
        await self.send_personal_message(
            {
                "type": "document_init",
                "document_id": document_id,
                "content": self.manager.get_document_content(document_id),
                "revision": self.manager.get_document_revision(document_id),
                "users": self.manager.get_online_users(document_id),
                "cursors": self.manager.get_user_cursors(document_id)
            },
            document_id,
            user_id
        )
        
        # 通知其他用户有新用户加入
        await self.broadcast_to_document(
            {
                "type": "user_joined",
                "document_id": document_id,
                "user_id": user_id,
                "user_info": user_info
            },
            document_id,
            exclude_user=user_id
        )
        
        logger.info(f"用户 {user_id} 连接到文档 {document_id}")
    
    def disconnect(self, document_id: str, user_id: str):
        """断开WebSocket连接"""
        if document_id in self.active_connections:
            if user_id in self.active_connections[document_id]:
                del self.active_connections[document_id][user_id]
                logger.info(f"用户 {user_id} 从文档 {document_id} 断开连接")
            
            # 如果文档中没有其他用户，清理连接
            if not self.active_connections[document_id]:
                del self.active_connections[document_id]
        
        # 移除在线用户
        self.manager.remove_user(document_id, user_id)
        
        # 通知其他用户有用户离开
        asyncio.create_task(
            self.broadcast_to_document(
                {
                    "type": "user_left",
                    "document_id": document_id,
                    "user_id": user_id
                },
                document_id
            )
        )
    
    async def send_personal_message(self, message: Dict[str, Any], document_id: str, user_id: str):
        """发送个人消息"""
        try:
            if (document_id in self.active_connections and 
                user_id in self.active_connections[document_id]):
                websocket = self.active_connections[document_id][user_id]
                await websocket.send_text(json.dumps(message, ensure_ascii=False))
        except Exception as e:
            logger.error(f"发送个人消息失败: {e}")
            self.disconnect(document_id, user_id)
    
    async def broadcast_to_document(self, message: Dict[str, Any], document_id: str, exclude_user: Optional[str] = None):
        """向文档中的所有用户广播消息"""
        if document_id not in self.active_connections:
            return
        
        disconnected_users = []
        
        for user_id, websocket in self.active_connections[document_id].items():
            # 跳过排除的用户
            if user_id == exclude_user:
                continue
                
            try:
                await websocket.send_text(json.dumps(message, ensure_ascii=False))
            except Exception as e:
                logger.error(f"向用户 {user_id} 发送消息失败: {e}")
                disconnected_users.append(user_id)
        
        # 清理断开连接的用户
        for user_id in disconnected_users:
            self.disconnect(document_id, user_id)
    
    async def handle_document_operation(self, message: Dict[str, Any], document_id: str, user_id: str):
        """处理文档操作"""
        try:
            operation_data = message.get("operation", {})
            operation = Operation.from_dict(operation_data)
            operation.user_id = user_id
            
            # 应用操作
            success = self.manager.apply_operation(document_id, operation)
            
            if success:
                # 广播操作给其他用户
                await self.broadcast_to_document(
                    {
                        "type": "document_operation",
                        "document_id": document_id,
                        "operation": operation.to_dict(),
                        "content": self.manager.get_document_content(document_id),
                        "revision": self.manager.get_document_revision(document_id)
                    },
                    document_id,
                    exclude_user=user_id
                )
            else:
                # 发送错误给操作用户
                await self.send_personal_message(
                    {
                        "type": "operation_error",
                        "document_id": document_id,
                        "error": "操作应用失败，请重新同步文档",
                        "current_revision": self.manager.get_document_revision(document_id)
                    },
                    document_id,
                    user_id
                )
        except Exception as e:
            logger.error(f"处理文档操作失败: {str(e)}")
            await self.send_personal_message(
                {
                    "type": "operation_error",
                    "document_id": document_id,
                    "error": f"操作处理失败: {str(e)}"
                },
                document_id,
                user_id
            )
    
    async def handle_cursor_update(self, message: Dict[str, Any], document_id: str, user_id: str):
        """处理光标更新"""
        cursor_data = message.get("cursor", {})
        self.manager.update_user_cursor(document_id, user_id, cursor_data)
        
        # 广播光标更新
        await self.broadcast_to_document(
            {
                "type": "cursor_update",
                "document_id": document_id,
                "user_id": user_id,
                "cursor": cursor_data
            },
            document_id,
            exclude_user=user_id
        )
    
    async def handle_lock_request(self, message: Dict[str, Any], document_id: str, user_id: str):
        """处理锁请求"""
        lock_info = message.get("lock_info", {})
        success = self.manager.acquire_lock(document_id, user_id, lock_info)
        
        if success:
            # 广播锁获取成功
            await self.broadcast_to_document(
                {
                    "type": "lock_acquired",
                    "document_id": document_id,
                    "user_id": user_id,
                    "lock_info": lock_info
                },
                document_id
            )
        else:
            # 发送锁被占用消息
            lock_info = self.manager.get_lock_info(document_id)
            await self.send_personal_message(
                {
                    "type": "lock_denied",
                    "document_id": document_id,
                    "lock_info": lock_info
                },
                document_id,
                user_id
            )
    
    async def handle_unlock_request(self, message: Dict[str, Any], document_id: str, user_id: str):
        """处理解锁请求"""
        self.manager.release_lock(document_id, user_id)
        
        # 广播锁释放
        await self.broadcast_to_document(
            {
                "type": "lock_released",
                "document_id": document_id,
                "user_id": user_id
            },
            document_id
        )
    
    async def handle_sync_request(self, message: Dict[str, Any], document_id: str, user_id: str):
        """处理同步请求"""
        # 发送文档当前状态
        await self.send_personal_message(
            {
                "type": "document_sync",
                "document_id": document_id,
                "content": self.manager.get_document_content(document_id),
                "revision": self.manager.get_document_revision(document_id),
                "operation_history": self.manager.get_operation_history(document_id)
            },
            document_id,
            user_id
        )


# 全局实例
enhanced_collaboration_service = EnhancedCollaborationService()


def get_enhanced_collaboration_service() -> EnhancedCollaborationService:
    """获取增强版协同编辑服务实例"""
    return enhanced_collaboration_service